[TRTLLM-11875][feat] BREAKING: MambaCacheManager based on KVCacheManagerV2 & agentic prefix caching by VALLIS-NERIA · Pull Request #16598 · NVIDIA/TensorRT-LLM

VALLIS-NERIA · 2026-07-20T06:04:56Z

Description

This PR adds the KV-cache-manager V2 implementation for hybrid Mamba models and moves Mamba state reuse onto explicit snapshot boundaries.

Key changes:

Add MambaHybridCacheManagerV2, including recurrent-state allocation, replay, promotion, dynamic-tree capacity, and deterministic stochastic-rounding state.
Add MambaStateConfig with periodic and fixed-position snapshot policies. The legacy mamba_state_cache_interval remains accepted only as a deprecated config-file alias.
Reuse Mamba state only at planned snapshot boundaries. Hybrid models that enable block reuse without any snapshot policy now warn and disable block reuse during model-aware config validation.
Add per-conversation reuse semantics: periodic Mamba snapshots are disabled and only explicit stable conversation boundaries are committed.
Add V2 Mamba state transfer to Python/NIXL disaggregated serving using explicit slot/layer strides.
Support one-model MTP and dynamic-tree MTP with V2, including target-tree mask regeneration and accepted-state promotion.
Handle attention-free hybrid PP stages in disaggregated transfer without zero-head duplication arithmetic.

This is an API/config migration: use kv_cache_config.mamba_state_config.periodic_snapshot_interval instead of kv_cache_config.mamba_state_cache_interval in Python code.

Test Coverage

Snapshot/config resolution tests, including V2-only fixed offsets and warning-plus-disable behavior when no snapshot policy exists.
V1/V2 hybrid manager, replay, snapshot planning, state promotion, scheduler, capacity, and per-conversation unit tests.
V1/V2 Mamba disaggregated extraction and transfer tests, including affine stride layouts and attention-free PP ranks.
Dynamic-tree metadata, relocation, capacity, and accepted-leaf state-promotion tests.
Parameterized Qwen3.5 and Nemotron Super V3 MTP accuracy tests are registered in the B200 CI lists for both V1 and V2.

Local validation after merging current main:

Relevant dynamic-tree/V2 tests: 4 passed.
Snapshot model-resolution tests: 6 passed.
Disaggregated peer/rank-info tests: 25 passed.
Qwen3.5 V2 MTP GSM8K: 94.8446 (threshold 91.327), passed.
All four new V1/V2 integration node IDs collect successfully.
Nemotron Super V3 V2 MTP: MMLU 85.38 (threshold 83.483) and GSM8K 92.15 (threshold 88.697), passed.
Pre-commit checks pass for all changed files.

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions).
If PR introduces API changes, an appropriate PR label is added, either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities.
CODEOWNERS is updated if ownership changes.
Documentation is updated as needed.
Update the TAVA architecture diagram if there is a significant design change.
The reviewers assigned automatically or manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, comment /bot help.

VALLIS-NERIA · 2026-07-20T06:07:48Z

/bot run

tensorrt-cicd · 2026-07-20T06:13:22Z

PR_Github #60284 [ run ] triggered by Bot. Commit: bfc9660 Link to invocation

tensorrt-cicd · 2026-07-20T06:31:52Z

PR_Github #60284 [ run ] completed with state FAILURE. Commit: bfc9660
/LLM/main/L0_MergeRequest_PR pipeline #48641 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

VALLIS-NERIA · 2026-07-20T15:18:45Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T15:24:18Z

PR_Github #60398 [ run ] triggered by Bot. Commit: 73e76f7 Link to invocation

tensorrt-cicd · 2026-07-20T18:43:39Z

PR_Github #60398 [ run ] completed with state FAILURE. Commit: 73e76f7
/LLM/main/L0_MergeRequest_PR pipeline #48738 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

nvpohanh · 2026-07-23T08:22:36Z

@lowsfer to review again.

tensorrt-cicd · 2026-07-23T08:38:24Z

PR_Github #61187 [ run ] completed with state SUCCESS. Commit: 9e6fbd5
/LLM/main/L0_MergeRequest_PR pipeline #49437 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-07-23T11:18:14Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T11:24:34Z

PR_Github #61282 [ run ] triggered by Bot. Commit: b060d39 Link to invocation

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

tensorrt-cicd · 2026-07-23T14:55:20Z

PR_Github #61282 [ run ] completed with state FAILURE. Commit: b060d39
/LLM/main/L0_MergeRequest_PR pipeline #49517 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

VALLIS-NERIA · 2026-07-23T14:56:42Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T15:03:16Z

PR_Github #61312 [ run ] triggered by Bot. Commit: d1ed15f Link to invocation

tensorrt-cicd · 2026-07-23T18:52:49Z

PR_Github #61312 [ run ] completed with state FAILURE. Commit: d1ed15f
/LLM/main/L0_MergeRequest_PR pipeline #49543 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

test_v2_hybrid_replay_bookkeeping_matches_checkpoint_predicate hardcoded a checkpoint-crossing scenario assuming replay_history_size == tokens_per_gen_step (5). Since replay_history_size became max(MIN_REPLAY_HISTORY_SIZE=16, tokens_per_gen_step), the scenario no longer crossed the checkpoint threshold, so prev_num_accepted_tokens accumulated (3+2=5) instead of resetting to 2. Derive the threshold from the manager's replay metadata and drive update_mamba_states until it actually crosses prev + step_width > history_size, so the test tracks the real checkpoint predicate regardless of the minimum history size. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-07-24T00:30:54Z

/bot run

tensorrt-cicd · 2026-07-24T00:36:25Z

PR_Github #61436 [ run ] triggered by Bot. Commit: 8f6c51f Link to invocation

tensorrt-cicd · 2026-07-24T00:56:32Z

PR_Github #61436 [ run ] completed with state FAILURE. Commit: 8f6c51f
/LLM/main/L0_MergeRequest_PR pipeline #49661 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

VALLIS-NERIA · 2026-07-24T01:02:00Z

/bot run

tensorrt-cicd · 2026-07-24T01:11:23Z

PR_Github #61447 [ run ] triggered by Bot. Commit: 8f6c51f Link to invocation

tensorrt-cicd · 2026-07-24T02:44:54Z

PR_Github #61447 [ run ] completed with state SUCCESS. Commit: 8f6c51f
/LLM/main/L0_MergeRequest_PR pipeline #49670 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Shixiaowei02

Could the auto resolver degrade back to False in dis-agg when the transceiver can't be PYTHON+NIXL? That avoids a startup hard-fail for existing UCX/MPI deployments. Besides, how far are we from making the default workflow? What tests still need to be completed?

lowsfer · 2026-07-24T03:48:55Z

+
+        if ssm_lc_id is not None:
+            page_ref = matched_blocks[-1].storage[ssm_lc_id]
+            if page_ref is None:


Seems duplicating previous code? we can move it to a local function and call it in both sites.

lowsfer · 2026-07-24T08:57:32Z

 class KVCacheDesc:
    capacity: int
    history_length: int
+    num_ssm_slots: int = 1


As we discussed, we shouldn't add this.

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-07-25T03:34:11Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-25T03:39:57Z

PR_Github #61677 [ run ] triggered by Bot. Commit: b0925e4 Link to invocation

tensorrt-cicd · 2026-07-25T07:21:42Z

PR_Github #61677 [ run ] completed with state FAILURE. Commit: b0925e4
/LLM/main/L0_MergeRequest_PR pipeline #49884 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

github-actions Bot assigned VALLIS-NERIA Jul 20, 2026

VALLIS-NERIA force-pushed the agent/v2-mamba-snapshot-reuse-core branch 2 times, most recently from 94db6dc to 6ffdf55 Compare July 20, 2026 10:19

VALLIS-NERIA added the api-breaking Accepted LLM API contract change that is backwards-incompatible label Jul 20, 2026

VALLIS-NERIA changed the title ~~[TRTLLM-11875][feat] MambaCacheManager based on KVCacheManagerV2 & agentic prefix caching~~ [TRTLLM-11875][feat] BREAKING: MambaCacheManager based on KVCacheManagerV2 & agentic prefix caching Jul 20, 2026

VALLIS-NERIA force-pushed the agent/v2-mamba-snapshot-reuse-core branch from 73e76f7 to f7770b5 Compare July 21, 2026 05:12

VALLIS-NERIA marked this pull request as ready for review July 21, 2026 05:14

VALLIS-NERIA requested review from a team as code owners July 21, 2026 05:14

VALLIS-NERIA requested review from QiJune, SimengLiu-nv, arysef, eopXD and pcastonguay July 21, 2026 05:14

qianlihuang mentioned this pull request Jul 23, 2026

[Exp][Hybrid KV Cache] Evaluate input-end vs decode-end checkpoint retention vllm-project/vllm#49574

Draft

Report SSM snapshot reuse statistics (By Agent)

b060d39

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA added 3 commits July 23, 2026 20:42

Use default V2 for MTP accuracy tests (By Agent)

6b529d4

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

[None][fix] Address pre-merge review findings (By Agent)

615d124

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

[Agent fix] Simplify V2 Mamba cache changes (By Agent)

d1ed15f

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

coderabbitai Bot mentioned this pull request Jul 23, 2026

[TRTLLM-14388][refactor] BREAKING: Remove 2-model code and tests #16380

Open

1 task

nvpohanh requested a review from lowsfer July 24, 2026 02:06

Shixiaowei02 approved these changes Jul 24, 2026

View reviewed changes

lowsfer reviewed Jul 24, 2026

View reviewed changes

VALLIS-NERIA added 3 commits July 24, 2026 22:45

Merge latest origin/main into V2 Mamba cache manager (By Agent)

fe0efc7

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

[None][fix] align V2 Mamba cache sizing (By Agent)

50520cb

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

Merge latest origin/main into V2 Mamba cache manager (By Agent)

b0925e4

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

Uh oh!

Conversation

VALLIS-NERIA commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

VALLIS-NERIA commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

VALLIS-NERIA commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

nvpohanh commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

VALLIS-NERIA commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

VALLIS-NERIA commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

VALLIS-NERIA commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

VALLIS-NERIA commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

Shixiaowei02 left a comment

Choose a reason for hiding this comment

Uh oh!

lowsfer Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

lowsfer Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

VALLIS-NERIA commented Jul 25, 2026

Uh oh!

tensorrt-cicd commented Jul 25, 2026

Uh oh!

tensorrt-cicd commented Jul 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

VALLIS-NERIA commented Jul 20, 2026 •

edited

Loading